from utils import *
import logging
import random
from collections import OrderedDict
import cv2
import imageio
from PIL import Image
import logging
import os
import re
import numpy as np
import sys
import random
import datetime
import torch
from utils import *
from torchvision.transforms import transforms
import torch.nn as nn
from networks import *

if __name__ == '__main__':
    modelpath = 'results/selfie2anime_params_latest.pt'
    input_nc = 3
    output_nc = 3
    ch = 32
    n_res = 4
    img_size = 256
    device = 'cpu'
    result_dir = 'test/A2B/results'
    input_dir = 'test/A2B/faces'
    genA2B = ResnetGenerator(input_nc=input_nc, output_nc=output_nc, ngf=ch, n_blocks=n_res, img_size=img_size, light=True).to(device)
    genB2A = ResnetGenerator(input_nc=input_nc, output_nc=output_nc, ngf=ch, n_blocks=n_res, img_size=img_size, light=True).to(device)
    params = torch.load(modelpath,map_location='cpu')

    ## 载入模型
    genA2B.load_state_dict(params['genA2B'])
    genB2A.load_state_dict(params['genB2A'])
    genA2B.eval()
    genB2A.eval()

    dummy_inputG = torch.randn((1, 3, 256, 256))
    torch.onnx.export(genA2B, dummy_inputG, "netA2B.onnx", verbose=True)

    if not os.path.isdir(result_dir):
        os.makedirs(result_dir)

    imagepaths = os.listdir(input_dir)
    transform = transforms.Compose([transforms.Resize((img_size, img_size)),transforms.ToTensor(),transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))])

    with torch.no_grad():
        for imagepath in imagepaths:
            image = Image.open(os.path.join(input_dir,imagepath))
            image = image.convert('RGB')
            image = transform(image)
            image.requires_grad = False
            image = image.unsqueeze(0).to(device)
            print(type(image))
            fake_A2B, _,_ = genA2B(image) ## 从A域到B域
            result = RGB2BGR(tensor2numpy(denorm(fake_A2B[0])))
            cv2.imwrite(os.path.join(result_dir,imagepath), result * 255.0)
